# -*- coding: utf-8 -*-

import pandas as pd
from sklearn.model_selection import StratifiedKFold, train_test_split


def split_data_with_label_counts(df, col, threshold):
    return df[df[col] > threshold], df[df[col] <= threshold]


if __name__ == "__main__":
    base_dir = './input'
    train_dir = base_dir + '/train_preliminary'
    train_click_log = pd.read_csv(train_dir + '/click_log.csv')
    for feat in train_click_log.columns:
        train_click_log[feat] = train_click_log[feat].apply(lambda x: None if x == '\\N' else x)
    grouped_cols = ['time', 'creative_id', 'click_times']
    train_click_log['label'] = train_click_log.groupby(grouped_cols).ngroup()
    label_counts = train_click_log['label'].value_counts()
    label_counts = label_counts.reset_index()
    label_counts.columns = ['label', 'label_counts']
    train_click_log = pd.merge(train_click_log, label_counts, on='label')

    n_splits = 3
    stratified_split_data, shuffle_split_data = split_data_with_label_counts(train_click_log,
                                                                             'label_counts', n_splits - 1)
    data_cols = ['time', 'user_id', 'creative_id', 'click_times']
    label_col = 'label'
    stratified_data = stratified_split_data[data_cols]
    stratified_label = stratified_split_data[label_col]

    shuffle_data = shuffle_split_data[data_cols]
    shuffle_label = shuffle_split_data[label_col]
    skf = StratifiedKFold(n_splits=n_splits)
    for i, (train_idx, valid_idx) in enumerate(skf.split(stratified_data, stratified_label)):
        stratified_train_data, stratified_valid_data = stratified_data.iloc[train_idx], stratified_data.iloc[valid_idx]

        shuffle_train_data, shuffle_valid_data, _, _ = train_test_split(shuffle_data, shuffle_label,
                                                                        test_size=1 / n_splits, random_state=i * 1024)

        train_data = pd.concat([stratified_train_data, shuffle_train_data])
        valid_data = pd.concat([stratified_valid_data, shuffle_valid_data])
        train_data.to_csv(base_dir + '/train_{}.csv'.format(str(i)))
        valid_data.to_csv(base_dir + '/valid_{}.csv'.format(str(i)))
